#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


###############################################################################
## Trifocal 2D camera location estimation from coplanar points. The
## input to this program are files produced by extract_edgels.


from __future__ import print_function

import sys

import time

import matplotlib
#matplotlib.use('Agg') 

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

from pylab import *

import scipy.io

from numpy import dot

from scipy.optimize import fmin_powell, fmin

from camori import EdgelSet, PicturePinhole, quaternion_to_matrix, measure_error, quaternion_product, dir_colors, color_distance

import simplejson

import Image

set_printoptions(precision=3)

def pplot(x,*a,**kw):
    if len(x.shape)>1:
        plot(x[:,0], x[:,1], *a, **kw)
    else:
        plot(x[0], x[1], *a, **kw)

def target_f(x, M, k):
    res = np.abs(dot(M[:,:-1], x) - M[:,-1])
    q = argsort(res)
    return res[q[k]]

def line_distance_worst(TR, pts):
    ABC = array([sin(TR[0]), cos(TR[0]), TR[1]])
    errs = np.abs(dot(pts, ABC[:2])+ABC[2])
    errs.sort()
    # return errs[errs.shape[0]*1/2]
    # return errs[-1]
    return errs[errs.shape[0]*4/5]

def line_distance_worst_2(pt, ABC):
    errs = np.abs(dot(ABC[:,:2], pt)+ABC[:,2])
    errs.sort()
    # return errs[(len(errs))/2]
    return errs[-3]
    # return errs[-1]

def estimate_line_H1(pts):
    ## First approximate by linear estimation.
    mupts = pts.mean(0)
    pts2 = c_[pts - mupts]
    
    AB = svd(pts2)[-1][-1]
    ABC = r_[-AB[1], AB[0], 0]
    TR = array([arctan2(ABC[1], ABC[0]), ABC[2]])
    
    TRopt = fmin_powell(line_distance_worst, TR, args=(pts2,), disp=False)
    ABCopt = array([sin(TRopt[0]), cos(TRopt[0]), TRopt[1]])
    ABCopt[2] -= dot(ABCopt[:2], mupts)
    return ABCopt

def estimate_vp_H1(ABC):
    ## First approximate by linear estimation.
    pt = svd(ABC)[-1][-1]
    pt = pt[:2] / pt[2]

    pt_opt = fmin(line_distance_worst_2, pt, args=(ABC,), disp=False)
    return pt_opt


if __name__ == '__main__':

    ## Plot stuff immediately
    ion()

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    #PlotStuff=False
    PlotStuff=True
    if sys.argv[0][-7:] == '-nox.py':
        PlotStuff = True

    

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print(sys.argv[0], '<job_file.json> <frame_number>')
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    oldgain = job_params['extraction_gain']
    gain = job_params['matching_gain']

    ## Get matrix for color distance comparison
    aa = np.load(fileroot +'/covs.npz')
    Mcov = aa['pcov']
    Sinv = inv(Mcov)

    ## Load edgels
    filename = fileroot+'/edgels/edgels-%04d.npz'%(int(sys.argv[2]))
    Ia = EdgelSet(filename, oldgain, gain)

    ## For each label, calculate the vanishing point direction.
    lines = [[],[],[]]
    for lab in range(3):
    
        Ia.calc_obs(lab)
        Ia.cluster_points(lab)

        heads = []
        dist=[]

        for ii in Ia.dr:
            end, distance = Ia.dr[ii]
            head, distance = Ia.dl[end]
            if head in heads:
                continue

            heads.append(head)
            dist.append(distance)

        heads = array(heads)
        dist = array(dist)

        Ntot = heads.shape[0]
        cut = Ntot/10
        if cut >= 4:
            heads = heads[dist.argsort()[-1:-cut:-1]]
            dist = dist[dist.argsort()[-1:-cut:-1]]
        else:
            heads = heads[dist.argsort()]
            dist = dist[dist.argsort()]

        for h in heads:
            x = h
            myl = []
            while x in Ia.dnr:
                myl.append(x)
                x = Ia.dnr[x]
            myl.append(x)

            pts = Ia.i_edgels[lab][myl,:2]
            lines[lab].append(pts)

            if PlotStuff:
                pplot(Ia.i_edgels[lab][myl,:2], '-+', color=dir_colors[lab])

    lineqs = [[],[],[]]
    for lab in range(3):
        for pts in lines[lab]:
            ## Solve to find each line equation coefficients.
            ABC = estimate_line_H1(pts)

            ## Append to list of lines from that VP
            lineqs[lab].append(ABC)

        lineqs[lab] = array(lineqs[lab])

    vps = zeros((3,2))
    for lab in range(3):
        vps[lab] = estimate_vp_H1(lineqs[lab])

    M = c_[vps[1]-vps[2],vps[0]-vps[2],vps[0]-vps[1]].T
    B = (vps*M).sum(1)

    ## The principal point
    pp = dot(pinv(M), B)

    ## Calculate the focal distance
    ivps = vps - pp
    fd = sqrt(-dot(ivps[1],ivps[2]))

    vv = c_[vps, fd*ones(3)].T
    vv = vv/c_[3*[sqrt((vv**2).sum(0))]]

    

    print('%d %f %f %f'%(int(sys.argv[2]), pp[0], pp[1], fd))

    if PlotStuff:
        ## Plot estimated lines
        for lab in range(3):
            ll = lineqs[lab]
            for ABC in ll:
                if np.abs(ABC[0]) < np.abs(ABC[1]):
                    p1 = array([-2000, -(-2000*ABC[0] + ABC[2])/ABC[1] ])
                    p2 = array([ 3000, -( 3000*ABC[0] + ABC[2])/ABC[1] ])
                else:
                    p1 = array([-(-2000*ABC[1] + ABC[2])/ABC[0],-2000 ])
                    p2 = array([-( 3000*ABC[1] + ABC[2])/ABC[0], 3000 ])

                plot([p1[0], p2[0]], [p1[1],p2[1]], '--', alpha=0.5, color=dir_colors[lab])

        plot([500], [375], 'ko')
        plot(pp[0], pp[1], 'bd')
        for lab in range(3):
            plot(vps[lab,0], vps[lab,1], 'd', color=dir_colors[lab])

        axis('equal')
        ylim(1250,-250)
        xlim(-500,1500)

    


